{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ab03a19a",
   "metadata": {},
   "source": [
    "# 注册算子"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e16b9bf4",
   "metadata": {},
   "source": [
    "Relay是TVM的高级IR，支持深度学习模型的表示、优化和执行。Relay的算子系统包括：\n",
    "- 算子元信息（名称、参数、属性等）\n",
    "- 类型关系（输入输出类型约束）\n",
    "- 计算函数（实现具体计算逻辑）\n",
    "- 策略函数（选择最佳实现）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9782960e",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5cf71110",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "\n",
    "def register(op_name, describe=\"\"):\n",
    "    \"\"\"Get the Op for a given name.\n",
    "    when the op_name is not registered, create a new empty op with the given name.\n",
    "    when the op_name has been registered, abort with an error message.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    op_name : str\n",
    "        The operator name\n",
    "\n",
    "    describe : Optional[str]\n",
    "        The operator description\n",
    "    \"\"\"\n",
    "\n",
    "    tvm.ir._ffi_api.RegisterOp(op_name, describe)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d3c57f5",
   "metadata": {},
   "source": [
    "## 1. 获取或创建算子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8878cce2",
   "metadata": {},
   "outputs": [],
   "source": [
    "op_name = \"my.operator\"\n",
    "try:\n",
    "    op = tvm.ir.Op.get(op_name)\n",
    "except:\n",
    "    op = register(op_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02f88ca2",
   "metadata": {},
   "source": [
    "## 2. 设置算子属性"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cea4e8f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "tvm.ir.Op.get(op_name).set_num_inputs(1)\n",
    "tvm.ir.Op.get(op_name).add_argument(\"data\", \"Tensor\", \"输入数据\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e945a752",
   "metadata": {},
   "source": [
    "## 3. 定义类型关系\n",
    "类型关系函数确保输入输出类型的一致性："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b92a99c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def my_op_type_rel(arg_types, attrs):\n",
    "    inputa_type = arg_types[0]\n",
    "    # print(inputa_type)\n",
    "    return tvm.relay.TensorType(inputa_type.shape, inputa_type.dtype)\n",
    "tvm.ir.Op.get(op_name).add_type_rel(\"MyOpTypeRel\", my_op_type_rel)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ff01a71",
   "metadata": {},
   "source": [
    "## 4. 设置计算函数\n",
    "\n",
    "计算函数实现具体的计算逻辑：\n",
    "```python\n",
    "@tvm.te.tag_scope(op_name)\n",
    "@tvm.relay.op.register_compute(op_name)\n",
    "def my_op_compute(attrs, inputs, output_type):\n",
    "    x = inputs[0]\n",
    "    out = x + 2 * 1\n",
    "    return [out]\n",
    "```\n",
    "\n",
    "或者："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5c4a4ae4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.my_op_compute(attrs, inputs, output_type)>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def my_op_compute(attrs, inputs, output_type):\n",
    "    x = inputs[0]\n",
    "    out = x + 2 * 1\n",
    "    return [out]\n",
    "\n",
    "tvm.relay.op.register_compute(op_name, my_op_compute, level=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ce943ff",
   "metadata": {},
   "source": [
    "## 5. 设置算子模式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3dde7b7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "tvm.ir.Op.get(op_name).set_attr(\"TOpPattern\", tvm.relay.op.OpPattern.ELEMWISE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52d821bb",
   "metadata": {},
   "source": [
    "## 6. 提供 Python API 接口\n",
    "\n",
    "为了方便使用，需要提供用户友好的Python API："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1bfc532e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def my_operator(data):\n",
    "    \"\"\"用户友好的API接口\"\"\"\n",
    "    op = tvm.ir.Op.get(op_name)\n",
    "    # attrs = tvm.ir.make_node(\"DictAttrs\")\n",
    "    return tvm.relay.Call(op, [data], attrs=None, type_args=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98dc5024",
   "metadata": {},
   "source": [
    "## 7. 实现算子调度\n",
    "\n",
    "为不同目标设备提供调度的实现策略："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bb2aae2e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GenericFunc(0x563dd948ff40)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@tvm.target.generic_func\n",
    "def schedule_special_op(attrs, outs, target):\n",
    "    with target:\n",
    "        outs = [outs] if isinstance(outs, tvm.te.tensor.Tensor) else outs\n",
    "        output = outs[0]\n",
    "        s = tvm.te.create_schedule(output.op)   \n",
    "        return s\n",
    "tvm.relay.op.op.register_schedule(op_name, schedule_special_op)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18119131",
   "metadata": {},
   "source": [
    "## 8. 测试与验证\n",
    "\n",
    "创建全面的测试用例确保算子功能正确："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "855f8141",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "测试通过!\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "# 创建测试数据\n",
    "np_data = np.random.randn(10).astype(\"float32\")\n",
    "\n",
    "# 创建Relay函数\n",
    "x = tvm.relay.var(\"x\", shape=(10,), dtype=\"float32\")\n",
    "y = my_operator(x)\n",
    "func = tvm.relay.Function([x], y)\n",
    "mod = tvm.IRModule.from_expr(func)\n",
    "\n",
    "# 执行计算\n",
    "intrp = tvm.relay.create_executor(kind=\"debug\", mod=mod, device=tvm.cpu(0))\n",
    "result = intrp.evaluate()(np_data)\n",
    "\n",
    "# 验证结果\n",
    "expected = np_data + 2 * 1   # 假设这是算子的预期行为\n",
    "np.testing.assert_allclose(result.numpy(), expected, rtol=1e-5)\n",
    "print(\"测试通过!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4c330f20",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">10</span>), float32]) {\n",
       "  my<span style=\"color: #A2F; font-weight: bold\">.</span>operator(<span style=\"color: #A2F; font-weight: bold\">%</span>x)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce86e031",
   "metadata": {},
   "source": [
    "## 创建统一接口"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "953c5dfc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RegisterOp(op_name='my.operator2', op=Op(my.operator2)) 测试通过!\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">10</span>), float32]) {\n",
       "  my<span style=\"color: #A2F; font-weight: bold\">.</span>operator2(<span style=\"color: #A2F; font-weight: bold\">%</span>x)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from dataclasses import dataclass, field, asdict\n",
    "import tvm\n",
    "\n",
    "def register_op(op_name, describe=\"\"):\n",
    "    \"\"\"Get the Op for a given name.\n",
    "    when the op_name is not registered, create a new empty op with the given name.\n",
    "    when the op_name has been registered, abort with an error message.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    op_name : str\n",
    "        The operator name\n",
    "\n",
    "    describe : Optional[str]\n",
    "        The operator description\n",
    "    \"\"\"\n",
    "\n",
    "    return tvm.ir._ffi_api.RegisterOp(op_name, describe)\n",
    "\n",
    "@tvm.target.generic_func\n",
    "def schedule_special_op(attrs, outs, target):\n",
    "    with target:\n",
    "        outs = [outs] if isinstance(outs, tvm.te.tensor.Tensor) else outs\n",
    "        output = outs[0]\n",
    "        s = tvm.te.create_schedule(output.op)   \n",
    "        return s\n",
    "\n",
    "@dataclass\n",
    "class RegisterOp:\n",
    "    op_name: str\n",
    "    op: tvm.ir.op.Op = field(init=False)\n",
    "\n",
    "    def __post_init__(self):\n",
    "        try:\n",
    "            self.op = tvm.ir.Op.get(self.op_name)\n",
    "        except:\n",
    "            register_op(self.op_name)\n",
    "            self.op = tvm.ir.Op.get(self.op_name)\n",
    "\n",
    "    def add_type_rel(self, rel_name, type_rel_func=None):\n",
    "        self.op.add_type_rel(rel_name, type_rel_func)\n",
    "\n",
    "    def register_compute(self, compute=None, level=10):\n",
    "        return tvm.relay.op.register_compute(self.op_name, compute, level=level)\n",
    "\n",
    "    def register_schedule(self):\n",
    "        return tvm.relay.op.op.register_schedule(self.op_name, schedule_special_op)\n",
    "\n",
    "    def register_pattern(self, pattern, level=10):\n",
    "        \"\"\"Register operator pattern for an op.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        op_name : str\n",
    "            The name of the op.\n",
    "\n",
    "        pattern : int\n",
    "            The pattern being used.\n",
    "\n",
    "        level : int\n",
    "            The priority level\n",
    "        \"\"\"\n",
    "        return tvm.ir.register_op_attr(self.op_name, \"TOpPattern\", pattern, level)\n",
    "\n",
    "    def __call__(self, args, attrs=None, type_args=None, span=None):\n",
    "        \"\"\"用户友好的回调 API 接口\"\"\"\n",
    "        return tvm.relay.Call(self.op, args, attrs=attrs, type_args=type_args, span=span)\n",
    "\n",
    "    def register(self, pattern, rel_name, type_rel_func=None, compute=None, level=10):\n",
    "        self.add_type_rel(rel_name, type_rel_func)\n",
    "        self.register_compute(compute, level=level)\n",
    "        self.register_schedule()\n",
    "        self.register_pattern(pattern, level=level)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    import numpy as np\n",
    "    # 创建测试数据\n",
    "    np_data = np.random.randn(10).astype(\"float32\")\n",
    "    op_name = \"my.operator2\"\n",
    "    def my_op_type_rel(arg_types, attrs):\n",
    "        inputa_type = arg_types[0]\n",
    "        return tvm.relay.TensorType(inputa_type.shape, inputa_type.dtype)\n",
    "    def my_op_compute(attrs, inputs, output_type):\n",
    "        x = inputs[0]\n",
    "        out = x + 2 * 1\n",
    "        return [out]\n",
    "    op = RegisterOp(op_name)\n",
    "    pattern = tvm.relay.op.OpPattern.ELEMWISE\n",
    "    rel_name = \"MyOpTypeRel\"\n",
    "    type_rel_func = my_op_type_rel\n",
    "    compute = my_op_compute\n",
    "    op.register(pattern, rel_name, type_rel_func, compute, level=10)\n",
    "    x = tvm.relay.var(\"x\", shape=(10,), dtype=\"float32\")\n",
    "    y = op([x])\n",
    "    mod = tvm.IRModule.from_expr(tvm.relay.Function([x], y))\n",
    "    # 执行计算\n",
    "    intrp = tvm.relay.create_executor(kind=\"vm\", mod=mod, device=tvm.cpu(0))\n",
    "    result = intrp.evaluate()(np_data)\n",
    "    # 验证结果\n",
    "    expected = np_data + 2 * 1   # 假设这是算子的预期行为\n",
    "    np.testing.assert_allclose(result.numpy(), expected, rtol=1e-5)\n",
    "    print(f\"{op} 测试通过!\")\n",
    "    mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30144834",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b3b1cc4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aix",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
